package io.yanglong.aiassistant;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentParser;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.ollama.OllamaEmbeddingModel;
import dev.langchain4j.store.embedding.redis.RedisEmbeddingStore;
import io.yanglong.aiassistant.config.Constant;

import java.net.URISyntaxException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

/**
 * RAG阶段一：将本地知识库向量化并存储到向量数据库
 */
public class MeiTuanRagLoader {

    public static void main(String[] args) throws URISyntaxException {
        final String regExp = "\\s*\\R\\s*\\R\\s*";
        //1.读取本地知识库文档
        Path path = Paths.get(Objects.requireNonNull(MeiTuanRagLoader.class.getClassLoader().getResource("meituan.txt")).toURI());
        DocumentParser documentParser = new TextDocumentParser();
        Document document = FileSystemDocumentLoader.loadDocument(path, documentParser);
        //2.对文件进行拆分，拆分为一个一个的知识条目
        DocumentSplitter documentSplitter = doc -> {
            List<TextSegment> segments = new ArrayList<>();
            String[] parts = doc.text().split(regExp);
            for (String part : parts) {
                System.out.println(part);
                segments.add(TextSegment.from(part));
            }
            return segments;
        };
        List<TextSegment> segments = documentSplitter.split(document);
        //3.对每个知识条目进行文本向量化并保存到向量数据库
        EmbeddingModel embeddingModel = OllamaEmbeddingModel.builder()
                .baseUrl("http://127.0.0.1:11434")
                .modelName(Constant.MODEL_NAME_DS_7B_QWEN)
                .build();
        RedisEmbeddingStore embeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                //维度，根据模型设置，ollama show [modelName]
                .dimension(Constant.DS_7B_QWEN_DIMENSION)
                .indexName("meituan-rag")
                .build();
        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();
        //向量化后的数据要和原文一起存入向量数据库
        embeddingStore.addAll(embeddings, segments);
    }
}
